# optimizer.py
"""
Description: 基于基本块 + DAG 的四元式优化（修正版）
Author: 天涯过客
Date: 2025-05-23 15:55:40
"""

from collections import namedtuple, defaultdict

Quad = namedtuple('Quad', ['op', 'arg1', 'arg2', 'res', 'param_type'])

PURE_OPS = {'+', '-', '*', '/', '%', '<', '>', '==', '!=', '&&', '||'}

def split_basic_blocks(quads):
    """
    划分基本块，边界为：
      * FUNC_BEGIN / FUNC_END
      * LABEL
      * 跳转指令 (JMP, JZ, JNZ) 之后一条
    返回 (block_index, block_quads) 列表
    """
    n = len(quads)
    leaders = {0}
    labels = {}
    for i,q in enumerate(quads):
        if q.op == 'LABEL':
            labels[q.res] = i
    for i,q in enumerate(quads):
        if q.op in ('JMP','JZ','JNZ'):
            # 跳转目标
            if q.res in labels:
                leaders.add(labels[q.res])
            # 跳转后下一条
            if i+1 < n:
                leaders.add(i+1)
        if q.op in ('FUNC_BEGIN','FUNC_END'):
            leaders.add(i)
    leaders = sorted(leaders)
    blocks = []
    for idx, start in enumerate(leaders):
        end = leaders[idx+1] if idx+1 < len(leaders) else n
        blocks.append((idx, quads[start:end]))
    return blocks

def optimize_block_dag(block):
    """
    对一个基本块内，识别连续的“纯运算子段”做 DAG 优化，其它指令原样保留。
    """
    res = []
    n = len(block)
    i = 0
    while i < n:
        # 如果是纯二元运算
        if block[i].op in PURE_OPS:
            j = i
            while j<n and block[j].op in PURE_OPS:
                j += 1
            segment = block[i:j]
            opt_segment = _dag_optimize(segment)
            res.extend(opt_segment)
            i = j
        else:
            res.append(block[i])
            i += 1
    return res

def _dag_optimize(segment):
    """
    针对一个纯运算四元式列表做 DAG 消除。
    返回优化后的四元式列表，自动过滤未被后续使用的临时写回。
    """
    # 1. 先扫描后续使用：哪些中间结果在段尾之后（剩余基本块）被使用？
    #    这里简化为：段内后续指令是否使用某临时。
    uses_after = defaultdict(bool)
    for q in segment[1:]:
        for operand in (q.arg1,q.arg2,q.res):
            uses_after[operand] = True

    node_map = {}           # key=(op,a1,a2) -> node_name
    node_seq = []           # 按出现顺序记录 (op,a1,a2,node)
    val_map = {}            # 原 res -> node
    counter = 0

    def new_node(op,a1,a2):
        nonlocal counter
        key = (op,a1,a2)
        if key in node_map:
            return node_map[key]
        name = f"t_opt{counter}"
        counter += 1
        node_map[key] = name
        node_seq.append((op,a1,a2,name))
        return name

    # 2. 构建 DAG
    for q in segment:
        # 左右操作数先映射
        a1 = val_map.get(q.arg1, q.arg1)
        a2 = val_map.get(q.arg2, q.arg2)
        # 新节点
        node = new_node(q.op,a1,a2)
        val_map[q.res] = node

    # 3. 输出：先所有节点计算
    out = [ Quad(op,a1,a2,node,None) for op,a1,a2,node in node_seq ]

    # 4. 写回：仅当原变量后续还会被使用
    for orig,node in val_map.items():
        if orig != node and uses_after.get(orig,False):
            out.append( Quad('ASSIGN', node, None, orig, None) )

    return out

def run_code_optimization(quads):
    """
    统一接口：返回 (flat_opt, bb_opt)
      - flat_opt: 仅优化后四元式
      - bb_opt: 在每块前加上 BB_START 标记的列表
    """
    blocks = split_basic_blocks(quads)
    flat, marked = [], []
    for idx, block in blocks:
        # 标记基本块开始
        marked.append( Quad('BB_START',None,None,f'block_{idx}',None) )
        opt = optimize_block_dag(block)
        flat.extend(opt)
        marked.extend(opt)
    return flat, marked
